module mpi_storage

use storage
use mpi
use mpi_basic_var

implicit none

logical ::  reorder =.false.
integer :: istatus(MPI_STATUS_SIZE)
integer :: n_front,n_back,n_left,n_right,n_top,n_bottom
integer, dimension(3) :: n_block,n_block_appo,n_cart,n_truebc_start,n_truebc_end
integer, dimension(:,:,:), allocatable :: cart_to_rank
integer, allocatable, dimension(:) :: all_st_x,all_en_x,all_sz_x
integer, allocatable, dimension(:) :: all_st_y,all_en_y,all_sz_y
integer, allocatable, dimension(:) :: all_st_z,all_en_z,all_sz_z
integer :: mpi_istart_x,mpi_iend_x
integer :: mpi_istart_y,mpi_iend_y
integer :: mpi_istart_z,mpi_iend_z
integer :: i_dim
integer :: mpi_comm_cart
integer :: ort_dir_1, ort_dir_2
integer :: n_nodes
integer :: i,j,k,cart_less,cart_more

end module mpi_storage

subroutine decompose_mpi

use mpi
use storage
use mpi_storage

implicit none

integer i_x,i_y,i_z

! Decompose procs in n_decomp rank blocks
n_block=1
n_cart=0
!-------------------------------------------------------------------------
!  OLD
!-------------------------------------------------------------------------
! call mpi_dims_create(n_proc,n_decomp,n_block,ierr)
! call mpi_cart_create(mpi_comm_world,n_decomp,n_block,periodicbc,  &
!                     reorder,mpi_comm_cart,ierr)
!print*,'MPI grid: ',n_block(1:n_decomp)
!! Obtaining process ids with in the cartesian grid
! call mpi_cart_coords(mpi_comm_cart,n_rank,n_decomp,n_cart,ierr)
!     call mpi_cart_rank(mpi_comm_cart, ncoords, newrank, ierr)
! call mpi_cart_shift(mpi_comm_cart, 0, 1, n_left   , n_right  , ierr)
! call mpi_cart_shift(mpi_comm_cart, 1, 1, n_bottom , n_top    , ierr)
! call mpi_cart_shift(mpi_comm_cart, 2, 1, n_back   , n_front  , ierr)
!-------------------------------------------------------------------------
!-------------------------------------------------------------------------
!  NEW
!-------------------------------------------------------------------------
n_nodes = n_proc / NCOREXNODE
! Decompone i nodi in griglia cartesiana 2D (o in generale (n_decomp-1)D)
if(n_decomp.eq.1) then
   n_block(1) = n_proc
   n_cart(1) = n_rank
else
   call mpi_dims_create(n_nodes,n_decomp-1,n_block,ierr)
   
   n_cart(node_dir)  = mod(n_rank,NCOREXNODE)
   
   ort_dir_1 = mod(node_dir-1,n_decomp)
   if(ort_dir_1.eq.0) ort_dir_1 = n_decomp
   
   if(n_decomp.eq.2) then
      n_cart(ort_dir_1) = n_rank/NCOREXNODE
   endif
   
   if(n_decomp.eq.3) then
      ort_dir_2 = mod(node_dir+1,n_decomp)
      if(ort_dir_2.eq.0) ort_dir_2 = n_decomp
      n_cart(ort_dir_1) = mod((n_rank/NCOREXNODE),n_block(1))
      n_cart(ort_dir_2) = (n_rank/NCOREXNODE)/n_block(1)
   !   n_cart(ort_dir_1) = mod((n_rank/NCOREXNODE),n_block(1))
   !   n_cart(ort_dir_2) = (n_rank/NCOREXNODE)/n_block(1)
   endif
   ! Rinomina decomposizione blocchi in 3D
   n_block_appo = n_block
   n_block(node_dir) = NCOREXNODE
   n_block(ort_dir_1) = n_block_appo(1)
   if(n_decomp.eq.3) n_block(ort_dir_2) = n_block_appo(2)
   if(n_rank .eq. 0) write(*,*) 'node_dir,ort_dir_1,ort_dir_2: ',node_dir,ort_dir_1,ort_dir_2
endif

! Calcola vettore (coordinate ---> rank)
if(n_rank .eq. 0) write(*,*) 'n_block:  ',n_block

if(n_decomp.eq.1) then
   n_cart(2:3) = 1
   allocate(cart_to_rank(0:n_block(1)-1,0:0,0:0))
   do i=0,n_block(1)-1
      cart_to_rank(i,0,0) = i
   enddo
endif
if(n_decomp.eq.2) then
   k = 1
   n_cart(3) = 1
   allocate(cart_to_rank(0:n_block(1)-1,0:n_block(2)-1,0))
   do i=0,n_block(node_dir)-1
   do j=0,n_block(ort_dir_1)-1
      if(node_dir.eq.1) cart_to_rank(i,j,0:0) = i + j*NCOREXNODE
      if(node_dir.eq.2) cart_to_rank(j,i,0:0) = i + j*NCOREXNODE
   enddo
   enddo
endif
if(n_decomp.eq.3) then
   allocate(cart_to_rank(0:n_block(1)-1,0:n_block(2)-1,0:n_block(3)-1))
   do i=0,n_block(node_dir)-1
   do j=0,n_block(ort_dir_1)-1
   do k=0,n_block(ort_dir_2)-1
      if(node_dir.eq.1) cart_to_rank(i,k,j) = i + j*NCOREXNODE + k*NCOREXNODE*n_block(ort_dir_1)
      if(node_dir.eq.2) cart_to_rank(j,i,k) = i + j*NCOREXNODE + k*NCOREXNODE*n_block(ort_dir_1)
      if(node_dir.eq.3) cart_to_rank(k,j,i) = i + j*NCOREXNODE + k*NCOREXNODE*n_block(ort_dir_1)
   enddo
   enddo
   enddo
endif

! Trova i confinanti (attenzione alle periodicita)
do i=1,n_decomp
   cart_less = n_cart(i)-1
   if(cart_less .eq. -1) then 
      if(periodicbc(i)) then
         cart_less = n_block(i)-1
      else
         cart_less = -3
      endif
   endif
   cart_more = n_cart(i)+1
   if(cart_more .eq. n_block(i)) then 
      if(periodicbc(i)) then
         cart_more = 0
      else
         cart_more = -3
      endif
   endif
   if(i.eq.1) then
      if(cart_less .eq. -3) then
        n_left = MPI_PROC_NULL
      else
        n_left = cart_to_rank(cart_less,n_cart(2),n_cart(3))
      endif
      if(cart_more .eq. -3) then
        n_right = MPI_PROC_NULL
      else
        n_right = cart_to_rank(cart_more,n_cart(2),n_cart(3))
      endif
!      write(100+n_rank,*) 'left:',cart_less,n_cart(2),n_cart(3),n_left
!      write(100+n_rank,*) 'right:',cart_more,n_cart(2),n_cart(3),n_right
   endif
   if(i.eq.2) then
      if(cart_less .eq. -3) then
        n_bottom = MPI_PROC_NULL
      else
        n_bottom = cart_to_rank(n_cart(1),cart_less,n_cart(3))
      endif
      if(cart_more .eq. -3) then
        n_top = MPI_PROC_NULL
      else
        n_top = cart_to_rank(n_cart(1),cart_more,n_cart(3))
      endif
!      write(100+n_rank,*) 'bottom:',n_cart(1),cart_less,n_cart(3),n_bottom
!      write(100+n_rank,*) 'top:',n_cart(1),cart_more,n_cart(3),n_top
   endif
   if(i.eq.3) then
      if(cart_less .eq. -3) then
        n_back = MPI_PROC_NULL
      else
        n_back = cart_to_rank(n_cart(1),n_cart(2),cart_less)
      endif
      if(cart_more .eq. -3) then
        n_front = MPI_PROC_NULL
      else
        n_front = cart_to_rank(n_cart(1),n_cart(2),cart_more)
      endif
!      write(100+n_rank,*) 'back:',n_cart(1),n_cart(2),cart_less,n_back
!      write(100+n_rank,*) 'front:',n_cart(1),n_cart(2),cart_more,n_front
   endif
enddo

if(n_decomp .eq. 1) then
!   write(100+n_rank,*) 'n_rank,n_cart(1),n_left,n_right: '
!   write(100+n_rank,*) n_rank,n_cart(1),n_left,n_right
endif
if(n_decomp .eq. 2) then
!   write(100+n_rank,*) 'n_rank,n_cart(1),n_cart(2),n_left,n_right,n_bottom,n_top: '
!   write(100+n_rank,*) n_rank,n_cart(1),n_cart(2),n_left,n_right,n_bottom,n_top
endif
if(n_decomp .eq. 3) then
!   write(100+n_rank,*) 'n_rank,n_cart(1),n_cart(2),n_cart(3),n_left,n_right,n_bottom,n_top,n_back,n_front: '
!   write(100+n_rank,*) n_rank,n_cart(1),n_cart(2),n_cart(3),n_left,n_right,n_bottom,n_top,n_back,n_front
endif

!print*,'MPI grid: ',n_block(1:n_decomp)
!print*,'RANK: ',n_rank, ' cart:', n_cart
!call MPI_FINALIZE(ierr)
!STOP
! Obtaining process ids with in the cartesian grid
!-------------------------------------------------------------------------
! For any n_decomp direction decompose nodes between blocks e set true_bc array
allocate(all_st_x(0:n_block(1)-1),all_en_x(0:n_block(1)-1),all_sz_x(0:n_block(1)-1))
  if((node_dir.eq.1).and.(NGPUXNODE .gt. 0)) then
     call MapDataToProc_fast(nt_x,n_block(1),NGPUXNODE,ratio_GPU,all_st_x,all_en_x,all_sz_x)
  else
     call MapDataToProc(nt_x,n_block(1),all_st_x,all_en_x,all_sz_x)
  endif
!  call MapDataToProc(nt_x,n_block(1),all_st,all_en,all_sz,ratio)
  mpi_istart_x = all_st_x(n_cart(1))   ;   mpi_iend_x = all_en_x(n_cart(1))
allocate(all_st_y(0:n_block(2)-1),all_en_y(0:n_block(2)-1),all_sz_y(0:n_block(2)-1))
all_st_y=1 ; all_en_y=1 ; all_sz_y=1
if(n_decomp .ge.2) then
     if((node_dir.eq.2).and.(NGPUXNODE .gt. 0)) then
        call MapDataToProc_fast(nt_y,n_block(2),NGPUXNODE,ratio_GPU,all_st_y,all_en_y,all_sz_y)
     else
        call MapDataToProc(nt_y,n_block(2),all_st_y,all_en_y,all_sz_y)
     endif
     mpi_istart_y = all_st_y(n_cart(2))   ;   mpi_iend_y = all_en_y(n_cart(2))
   endif
allocate(all_st_z(0:n_block(3)-1),all_en_z(0:n_block(3)-1),all_sz_z(0:n_block(3)-1))
all_st_z=1 ; all_en_z=1 ; all_sz_z=1
if(n_decomp .eq.3) then
   if((node_dir.eq.3) .and. (NGPUXNODE .gt. 0)) then
      call MapDataToProc_fast(nt_z,n_block(3),NGPUXNODE,ratio_GPU,all_st_z,all_en_z,all_sz_z)
   else
       call MapDataToProc(nt_z,n_block(3),all_st_z,all_en_z,all_sz_z)
   endif
   mpi_istart_z = all_st_z(n_cart(3))   ;   mpi_iend_z = all_en_z(n_cart(3))
endif

do i_dim = 1,n_decomp
   n_truebc_start(i_dim) = 1
   n_truebc_end(i_dim) = 1
   if(n_cart(i_dim) .ne. 0) n_truebc_start(i_dim) = 0
   if(n_cart(i_dim) .ne. n_block(i_dim)-1) n_truebc_end(i_dim) = 0
enddo

! correzione per condizioni periodiche che sono sempre false condizoni. controllare!!!???????
do i_dim = 1,n_decomp
   if(periodicbc(i_dim)) then
      n_truebc_start(i_dim) = 0
      n_truebc_end(i_dim) = 0
   endif
enddo

print*,'n_rank:',n_rank,'st: ',n_truebc_start,' en: ',n_truebc_end
print*,'Rank: ',n_rank,'; n_cart: ',n_cart(1:n_decomp)
print*,'Rank: ',n_rank,'; neigh: ',n_front,n_back,n_top,n_bottom,n_left,n_right
istart_x = mpi_istart_x
iend_x   = mpi_iend_x
print*,'mpi_istart_x,mpi_iend_x: ',mpi_istart_x,mpi_iend_x
if(n_decomp .ge.2) then
   istart_y = mpi_istart_y
   iend_y   = mpi_iend_y
endif
print*,'mpi_istart_y,mpi_iend_y: ',mpi_istart_y,mpi_iend_y
if(n_decomp .eq.3) then
   istart_z = mpi_istart_z
   iend_z   = mpi_iend_z
   print*,'mpi_istart_z,mpi_iend_z: ',mpi_istart_z,mpi_iend_z
endif

if(n_decomp .eq.1) then
   isize_x = iend_x-istart_x+1+2*ng_x
   isize_max_1 = isize_x
   ran_nvec = isize_max_1
   if(isize_x < ng_x) then
     print*,'rank:',n_rank,': too many procs along x! isize_x:',isize_x,'< ng_x:',ng_x
     call MPI_FINALIZE(ierr)
   end if
endif
if(n_decomp .eq.2) then
   isize_x = iend_x-istart_x+1+2*ng_x
   isize_y = iend_y-istart_y+1+2*ng_y
   isize_max_1 = max(isize_x,isize_y)
   if(isize_max_1 == isize_x) isize_max_2 = isize_y
   if(isize_max_1 == isize_y) isize_max_2 = isize_x
   ran_nvec = isize_max_1 * isize_max_2
   if(isize_y < ng_y) then
     print*,'rank:',n_rank,': too many procs along y! isize_y:',isize_y,'< ng_y:',ng_y
     call MPI_FINALIZE(ierr)
   end if
endif
if(n_decomp .eq.3) then
   isize_x = iend_x-istart_x+1+2*ng_x
   isize_y = iend_y-istart_y+1+2*ng_y
   isize_z = iend_z-istart_z+1+2*ng_z
   isize_max_1 = max(isize_x,isize_y,isize_z)
   if(isize_max_1 == isize_x) isize_max_2 = max(isize_y,isize_z)
   if(isize_max_1 == isize_y) isize_max_2 = max(isize_x,isize_z)
   if(isize_max_1 == isize_z) isize_max_2 = max(isize_x,isize_y)
   ran_nvec = isize_max_1 * isize_max_2
   if(isize_z < ng_z) then
     print*,'rank:',n_rank,': too many procs along z! isize_z:',isize_z,'< ng_z:',ng_z
     call MPI_FINALIZE(ierr)
   end if
endif

!call MPI_FINALIZE(ierr)
!STOP 'test_MPI'

end subroutine decompose_mpi
